import torch
import jieba
import re
import csv

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

n_step = 1000

word2idx = {}
idx2word = {}

with open("./data/word2idx.csv", "r") as f:
    reader = csv.reader(f)
    for row in reader:
        key, value = row
        word2idx[key] = int(value)  # 如果需要，可以将值转换为适当类型

with open("./data/idx2word.csv", "r") as f:
    reader = csv.reader(f)
    for row in reader:
        key, value = row
        idx2word[int(key)] = value  # 如果需要，可以将值转换为适当类型


def preprocess_sentence(sentence, word2idx, n_step):
    # 字符过滤
    sentence = re.sub(r'[^\u4e00-\u9fa5a-zA-Z0-9，。！？]', '', sentence)
    # 分词
    words = jieba.cut(sentence)
    # 索引映射
    array = [2] + [word2idx.get(w, word2idx["<UNK>"]) for w in words] + [3]

    # 序列填充
    if len(array) < n_step:
        array += [0] * (n_step - len(array))
    return array


def make_data(input_sentence, output_sentence):
    input_array, output_array = [], []
    for input_sent, output_sent in zip(input_sentence, output_sentence):
        input_array.append(preprocess_sentence(input_sent, word2idx, n_step))
        output_array.append(preprocess_sentence(output_sent, word2idx, n_step))

    return torch.Tensor(input_array).long(), torch.Tensor(output_array).long()


# 加载 TorchScript 模型
model = torch.jit.load("pytorch-transformer.pt")
model.eval()  # 将模型设置为评估模式


def predictSay(text):
    model.eval()
    myinput, _ = make_data([text], [""])

    with torch.no_grad():
        tgt = torch.tensor([[word2idx["<SOS>"]]]).to(device)
        decoded = []

        # 逐步生成输出序列
        for _ in range(n_step):
            output = model(myinput.to(device), tgt)  # 前向传播
            next_word_idx = output[:, -1, :].argmax(dim=-1).item()  # 获取下一个单词的索引
            decoded.append(idx2word[next_word_idx])  # 将索引转换为单词并添加到结果中
            # 如果生成了 <EOS>，停止生成
            if next_word_idx == word2idx["<EOS>"]:
                break

            # 更新目标序列
            tgt = torch.cat([tgt, torch.tensor([[next_word_idx]], dtype=torch.long).to(device)], dim=1)

    return ''.join(decoded).replace("<SOS>", "").replace("<PAD>", "").replace("<EOS>", "")


print(predictSay("你多大了"))
